#!/home/agilex/miniconda3/envs/graspnet/bin/python
# -- coding: UTF-8
import rospy
import std_msgs.msg
from tf2_ros import TransformBroadcaster
from geometry_msgs.msg import TransformStamped
import numpy as np


import struct
import std_msgs
from sensor_msgs.msg import PointCloud2, PointField
import sensor_msgs.point_cloud2 as pc2

import pyrealsense2 as rs

import queue
import threading
import open3d as o3d

intrinsics_matrix = np.array([[609.0, 0., 320.],
                         [ 0.,609.0, 240],
                         [0,0,1]], dtype=np.float32)

factor_depth = np.array([[1000.]])

class CameraInfo():
    def __init__(self, width, height, fx, fy, cx, cy, scale):
        self.width = width
        self.height = height
        self.fx = fx
        self.fy = fy
        self.cx = cx
        self.cy = cy
        self.scale = scale

def rad2deg(rad):
    return rad / np.pi * 180. 


def rgbd2cloud(color, depth, camera, rgb_flag = False, num_point = 10000):
    assert(depth.shape[0] == camera.height and depth.shape[1] == camera.width)
    xmap = np.arange(camera.width)
    ymap = np.arange(camera.height)
    xmap, ymap = np.meshgrid(xmap, ymap)  # 创建图像网格
    points_z = depth / camera.scale       # 深度信息  使用了广播机制
   
    points_x = (xmap - camera.cx) * points_z / camera.fx  # X = (u - cx) * Z /fx
    points_y = (ymap - camera.cy) * points_z / camera.fy  # Y = (v - cx) * Y /fy
    cloud = np.stack([points_x, points_y, points_z], axis=-1)  # 沿着最后一个轴增加维度
    if rgb_flag:
        color = color[:, :, ::-1]   # bgr2rgb
        cloud = np.concatenate((cloud, color), axis=2)
        # print("cloud: ", cloud.shape)
        cloud = cloud.reshape([-1, 6])
    else:
        cloud = cloud.reshape([-1, 3])
    cloud = cloud[(cloud[:, 2] > 0.2) & (cloud[:, 2] < 2)]
    
    cloud_size = cloud.shape[0]
    if(cloud_size >= num_point):
        idxs = np.random.choice(cloud.shape[0] , num_point, replace=False)
    else:
        idxs1 = np.arange(cloud_size)  # 补点
        idxs2 = np.random.choice(cloud_size, num_point - cloud_size, replace=True)
        idxs = np.concatenate([idxs1, idxs2], axis=0)
        idxs = np.random.choice(cloud.shape[0] , num_point, replace=False)
    cloud = cloud[idxs]
    
    return cloud


def open3d2msg(open3d_cloud, frame_id="map"):
    points = np.asarray(open3d_cloud.points, dtype=np.float32)
    colors = (np.asarray(open3d_cloud.colors) * 255).astype(np.uint8) if len(open3d_cloud.colors) > 0 else np.full((points.shape[0], 3), 255, dtype=np.uint8)

    # Prepare data for PointCloud2
    cloud_data = np.zeros((points.shape[0], 4), dtype=np.float32)
    cloud_data[:, 0:3] = points

    for i in range(colors.shape[0]):
        r, g, b = colors[i]
        rgb_packed = (r << 16) | (g << 8) | b
        rgb_float = struct.unpack('f', struct.pack('I', rgb_packed))[0]
        cloud_data[i, 3] = rgb_float

    # Flatten data to byte array
    cloud_data = cloud_data.flatten().tobytes()

    # Create PointCloud2 message
    header = std_msgs.msg.Header()
    header.stamp = rospy.Time.now()
    header.frame_id = frame_id

    fields = [
        PointField('x', 0, PointField.FLOAT32, 1),
        PointField('y', 4, PointField.FLOAT32, 1),
        PointField('z', 8, PointField.FLOAT32, 1),
        PointField('rgb', 12, PointField.FLOAT32, 1),
    ]

    point_cloud_msg = PointCloud2(
        header=header,
        height=1,
        width=points.shape[0],
        is_dense=False,
        is_bigendian=False,
        fields=fields,
        point_step=16,
        row_step=16 * points.shape[0],
        data=cloud_data
    )
    return point_cloud_msg


def convert_numpy_to_pointcloud2(np_array):
    """
    Convert numpy array to sensor_msgs/PointCloud2.
    
    np_array: A numpy array of shape (n, 6) with each row as (x, y, z, r, g, b)
    
    Returns:
    PointCloud2 message
    """
    header = std_msgs.msg.Header()
    header.stamp = rospy.Time.now()
    header.frame_id = "map"

    fields = [
        PointField('x', 0, PointField.FLOAT32, 1),
        PointField('y', 4, PointField.FLOAT32, 1),
        PointField('z', 8, PointField.FLOAT32, 1),
        PointField('rgba', 12, PointField.UINT32, 1),
    ]

    point_cloud = []
    for point in np_array:
        x, y, z, r, g, b = point
        r = min(max(int(r), 0), 255)
        g = min(max(int(g), 0), 255)
        b = min(max(int(b), 0), 255)
        a = 255  # Full opacity
        rgba = struct.unpack('I', struct.pack('BBBB', b, g, r, a))[0]
        point_cloud.append([x, y, z, rgba])

    pc2_msg = pc2.create_cloud(header, fields, point_cloud)
    
    return pc2_msg


class PublishNode:
    def __init__(self) -> None:
        rospy.init_node("grasp_node")
        self.tf_broad = TransformBroadcaster()
        self.msg_tf = TransformStamped()
        self.rate = rospy.Rate(10)
        self.pub_cloud = rospy.Publisher("/points", PointCloud2, queue_size=10)
        self.camera = CameraInfo(640, 480, intrinsics_matrix[0][0], intrinsics_matrix[1][1], intrinsics_matrix[0][2], intrinsics_matrix[1][2], factor_depth)
        
        # 初始化相机
        self.pipeline = None
        self.align = None
        self.init_camera()

        self.cloudQue = queue.Queue()
        self.cloud = None
        self.num_points = 10000

        t1 = threading.Thread(target=self.produce_rgbd)
        t1.start()
        
        
        while not rospy.is_shutdown():
            # t2 = None
            if self.cloudQue.empty():
                self.rate.sleep()
                continue
            self.cloud = self.cloudQue.get()
            print("main_function : ", self.cloud.shape)
            
            # pcd = o3d.geometry.PointCloud()
            # pcd.points = o3d.utility.Vector3dVector(self.cloud[:, :3].astype(np.float32))
            # pcd.colors = o3d.utility.Vector3dVector(np.array(self.cloud[:, 3:], dtype=np.float32) / 255.0)

            t2 = threading.Thread(target=self.publish_cloud, args=[self.cloud])
            t2.start()
            
            print(self.cloudQue.qsize())
            input("press any key:")
            self.rate.sleep()

    def produce_rgbd(self):
        while not rospy.is_shutdown():
            try:
                # 等待深度数据帧和RGB数据帧，设置等待时间为10秒
                frames = self.pipeline.wait_for_frames(timeout_ms=10000)
                aligned_frames = self.align.process(frames)
                if not aligned_frames:
                    continue  # If alignment fails, go bac
                
                depth_frame = frames.get_depth_frame()
                color_frame = frames.get_color_frame()
                
                if not depth_frame or not color_frame:
                    continue
                # 获取深度图像的原始数据
                depth_data = np.asanyarray(depth_frame.get_data())
                # 获取RGB图像的原始数据
                color_data = np.asanyarray(color_frame.get_data())

                cloud = rgbd2cloud(color_data, depth_data, self.camera, True, self.num_points)
                self.cloudQue.put(cloud)
                if self.cloudQue.qsize() > 10:
                    self.cloudQue.get()
            except RuntimeError as e:
                print(f"等待帧时发生错误: {e}")

    def init_camera(self):
        # 初始化深度相机
        self.pipeline = rs.pipeline()
        config = rs.config()

        # 配置深度和颜色流
        config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
        config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)

        # 开始更新配置和同步设置
        self.pipeline.start(config)
        self.align = rs.align(rs.stream.color)  # Create align object for depth-color alignment


    def publish_tf(self, frame_id, child_frame_id, q, t):
        self.msg_tf.header.frame_id = frame_id
        self.msg_tf.child_frame_id = child_frame_id
        # -0.002 -0.015  0.201 -0.295 -0.293 -0.645  0.641
        self.msg_tf.transform.rotation.x = q[0]
        self.msg_tf.transform.rotation.y = q[1]
        self.msg_tf.transform.rotation.z = q[2]
        self.msg_tf.transform.rotation.w = q[3]
        self.msg_tf.transform.translation.x = t[0]
        self.msg_tf.transform.translation.y = t[1]
        self.msg_tf.transform.translation.z = t[2]
       
        self.msg_tf.header.stamp = rospy.Time.now()
        self.tf_broad.sendTransform(self.msg_tf)
            
    def publish_cloud(self, cloud):
        
        while not rospy.is_shutdown():
            # print(cloud)
            # msg_cloud = open3d2msg(cloud)
            msg_cloud = convert_numpy_to_pointcloud2(cloud)
            self.pub_cloud.publish(msg_cloud)
            self.rate.sleep()

if __name__ == "__main__":
    pub_node = PublishNode()